{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "200e725e",
   "metadata": {},
   "source": [
    "# Tutorial Notebook 9: Natural Language Interpretation of Single-Cell Datasets with a C2S Foundation Model\n",
    "\n",
    "In this tutorial notebook, we demonstrate how to generate natural language interpretations of single-cell datasets using a Cell2Sentence (C2S) model. We will utilize a pretrained C2S foundation model to generate descriptive summaries for sets of cells, treating the task as a form of natural language interpretation of complex biological data.\n",
    "\n",
    "This notebook will guide you through the following steps:\n",
    "1.  Load an immune tissue single-cell dataset from Domínguez Conde et al. (preprocessed in tutorial notebook 0).\n",
    "2.  Load a pretrained C2S model.\n",
    "3.  Format multi-cell inputs into prompts for natural language interpretation.\n",
    "4.  Use the C2S model to generate insightful summaries for different sets of cells."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c85ef67",
   "metadata": {},
   "source": [
    "We will begin by importing the necessary libraries. These include standard Python libraries for data manipulation, as well as specific modules from the `cell2sentence` package for handling single-cell data and C2S models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9b66f739",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "\n",
    "from datasets import Dataset\n",
    "import numpy as np\n",
    "import anndata\n",
    "import scanpy as sc\n",
    "import torch\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e9149a9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cell2sentence as cs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "44342a01",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 1234\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "894938d9",
   "metadata": {},
   "source": [
    "# Load Data\n",
    "\n",
    "Next, we will load the preprocessed single-cell dataset from tutorial 0. This dataset has been filtered and normalized, making it ready for conversion into cell sentences.\n",
    "\n",
    "<font color='red'>Please ensure you have run the preprocessing steps in Tutorial 0 before executing the following code if you are using your own dataset.</font> Make sure the file path in <font color='gold'>DATA_PATH</font> is correctly set to your preprocessed data file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b727f9d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 29773 × 23944\n",
       "    obs: 'cell_type', 'tissue', 'batch_condition', 'organism', 'assay', 'sex', 'n_genes', 'n_genes_by_counts', 'total_counts', 'total_counts_mt', 'pct_counts_mt'\n",
       "    var: 'gene_name', 'ensembl_id', 'n_cells', 'mt', 'n_cells_by_counts', 'mean_counts', 'pct_dropout_by_counts', 'total_counts'\n",
       "    uns: 'batch_condition_colors', 'cell_type_colors', 'log1p', 'neighbors', 'pca', 'tissue_colors', 'umap'\n",
       "    obsm: 'X_pca', 'X_umap'\n",
       "    varm: 'PCs'\n",
       "    obsp: 'connectivities', 'distances'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "DATA_PATH = \"/home/sr2464/scratch/C2S_API_Testing/Data/dominguez_conde_immune_tissue_two_donors_preprocessed_tutorial_0.h5ad\"\n",
    "adata = anndata.read_h5ad(DATA_PATH)\n",
    "adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0cb6535d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>cell_type</th>\n",
       "      <th>tissue</th>\n",
       "      <th>batch_condition</th>\n",
       "      <th>organism</th>\n",
       "      <th>assay</th>\n",
       "      <th>sex</th>\n",
       "      <th>n_genes</th>\n",
       "      <th>n_genes_by_counts</th>\n",
       "      <th>total_counts</th>\n",
       "      <th>total_counts_mt</th>\n",
       "      <th>pct_counts_mt</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Pan_T7935490_AAACCTGCAAATTGCC</th>\n",
       "      <td>CD4-positive helper T cell</td>\n",
       "      <td>ileum</td>\n",
       "      <td>A29</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>10x 5' v1</td>\n",
       "      <td>female</td>\n",
       "      <td>2191</td>\n",
       "      <td>2191</td>\n",
       "      <td>6542.0</td>\n",
       "      <td>327.0</td>\n",
       "      <td>4.998472</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Pan_T7935490_AAACGGGCATCTGGTA</th>\n",
       "      <td>CD8-positive, alpha-beta memory T cell</td>\n",
       "      <td>ileum</td>\n",
       "      <td>A29</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>10x 5' v1</td>\n",
       "      <td>female</td>\n",
       "      <td>2046</td>\n",
       "      <td>2046</td>\n",
       "      <td>5871.0</td>\n",
       "      <td>429.0</td>\n",
       "      <td>7.307103</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Pan_T7935490_AAACGGGTCTTGCATT</th>\n",
       "      <td>CD8-positive, alpha-beta memory T cell</td>\n",
       "      <td>ileum</td>\n",
       "      <td>A29</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>10x 5' v1</td>\n",
       "      <td>female</td>\n",
       "      <td>2129</td>\n",
       "      <td>2129</td>\n",
       "      <td>7248.0</td>\n",
       "      <td>337.0</td>\n",
       "      <td>4.649559</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Pan_T7935490_AAAGCAATCATCGCTC</th>\n",
       "      <td>CD8-positive, alpha-beta memory T cell</td>\n",
       "      <td>ileum</td>\n",
       "      <td>A29</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>10x 5' v1</td>\n",
       "      <td>female</td>\n",
       "      <td>1262</td>\n",
       "      <td>1262</td>\n",
       "      <td>3927.0</td>\n",
       "      <td>305.0</td>\n",
       "      <td>7.766743</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Pan_T7935490_AAAGTAGCAGTCACTA</th>\n",
       "      <td>gamma-delta T cell</td>\n",
       "      <td>ileum</td>\n",
       "      <td>A29</td>\n",
       "      <td>Homo sapiens</td>\n",
       "      <td>10x 5' v1</td>\n",
       "      <td>female</td>\n",
       "      <td>2248</td>\n",
       "      <td>2248</td>\n",
       "      <td>6574.0</td>\n",
       "      <td>1083.0</td>\n",
       "      <td>16.473988</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                            cell_type tissue  \\\n",
       "Pan_T7935490_AAACCTGCAAATTGCC              CD4-positive helper T cell  ileum   \n",
       "Pan_T7935490_AAACGGGCATCTGGTA  CD8-positive, alpha-beta memory T cell  ileum   \n",
       "Pan_T7935490_AAACGGGTCTTGCATT  CD8-positive, alpha-beta memory T cell  ileum   \n",
       "Pan_T7935490_AAAGCAATCATCGCTC  CD8-positive, alpha-beta memory T cell  ileum   \n",
       "Pan_T7935490_AAAGTAGCAGTCACTA                      gamma-delta T cell  ileum   \n",
       "\n",
       "                              batch_condition      organism      assay  \\\n",
       "Pan_T7935490_AAACCTGCAAATTGCC             A29  Homo sapiens  10x 5' v1   \n",
       "Pan_T7935490_AAACGGGCATCTGGTA             A29  Homo sapiens  10x 5' v1   \n",
       "Pan_T7935490_AAACGGGTCTTGCATT             A29  Homo sapiens  10x 5' v1   \n",
       "Pan_T7935490_AAAGCAATCATCGCTC             A29  Homo sapiens  10x 5' v1   \n",
       "Pan_T7935490_AAAGTAGCAGTCACTA             A29  Homo sapiens  10x 5' v1   \n",
       "\n",
       "                                  sex  n_genes  n_genes_by_counts  \\\n",
       "Pan_T7935490_AAACCTGCAAATTGCC  female     2191               2191   \n",
       "Pan_T7935490_AAACGGGCATCTGGTA  female     2046               2046   \n",
       "Pan_T7935490_AAACGGGTCTTGCATT  female     2129               2129   \n",
       "Pan_T7935490_AAAGCAATCATCGCTC  female     1262               1262   \n",
       "Pan_T7935490_AAAGTAGCAGTCACTA  female     2248               2248   \n",
       "\n",
       "                               total_counts  total_counts_mt  pct_counts_mt  \n",
       "Pan_T7935490_AAACCTGCAAATTGCC        6542.0            327.0       4.998472  \n",
       "Pan_T7935490_AAACGGGCATCTGGTA        5871.0            429.0       7.307103  \n",
       "Pan_T7935490_AAACGGGTCTTGCATT        7248.0            337.0       4.649559  \n",
       "Pan_T7935490_AAAGCAATCATCGCTC        3927.0            305.0       7.766743  \n",
       "Pan_T7935490_AAAGTAGCAGTCACTA        6574.0           1083.0      16.473988  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.obs.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "78b71e4d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 29773 × 23944\n",
       "    obs: 'cell_type', 'tissue', 'batch_condition', 'organism', 'sex'\n",
       "    var: 'gene_name', 'ensembl_id', 'n_cells', 'mt', 'n_cells_by_counts', 'mean_counts', 'pct_dropout_by_counts', 'total_counts'\n",
       "    uns: 'batch_condition_colors', 'cell_type_colors', 'log1p', 'neighbors', 'pca', 'tissue_colors', 'umap'\n",
       "    obsm: 'X_pca', 'X_umap'\n",
       "    varm: 'PCs'\n",
       "    obsp: 'connectivities', 'distances'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.obs = adata.obs[[\"cell_type\", \"tissue\", \"batch_condition\", \"organism\", \"sex\"]]\n",
    "adata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2be1277",
   "metadata": {},
   "source": [
    "# Cell2Sentence Conversion\n",
    "\n",
    "In this section, we will convert our AnnData object into a format that the C2S model can understand. This involves transforming the gene expression data of each cell into a \"cell sentence,\" which is a string of gene names ordered by their expression level. We will also prepare the vocabulary from our dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "842861e5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29773/29773 [00:09<00:00, 2995.86it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['cell_name', 'cell_sentence', 'cell_type', 'tissue', 'batch_condition', 'organism', 'sex'],\n",
       "    num_rows: 29773\n",
       "})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata_obs_cols_to_keep = adata.obs.columns.tolist()\n",
    "\n",
    "# Create Arrow dataset and vocabulary from AnnData\n",
    "arrow_ds, vocabulary = cs.CSData.adata_to_arrow(\n",
    "    adata=adata,\n",
    "    random_state=SEED,\n",
    "    sentence_delimiter=' ',\n",
    "    label_col_names=adata_obs_cols_to_keep\n",
    ")\n",
    "arrow_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ac96a758",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'cell_name': 'Pan_T7935490_AAACCTGCAAATTGCC',\n",
       " 'cell_sentence': 'RPLP1 ACTB EEF1A1 HSP90AA1 TMSB4X B2M FTH1 KLF6 HSPA1B MALAT1 RPS12 HSPA8 RPL13 MT-CO1 ATF3 MT-CO2 RPL41 TPT1 MT-CO3 RPS19 HLA-B RPL10 RPS4X RPL28 MT-CYB DUSP1 RPL30 MT-ND4L RPS15 FOS RPL34 RPS2 RPLP2 MT-ND3 RPS18 RPS8 TRBV7-2 RPL32 RPS3 ANXA1 RPL11 HLA-C RPS27 ACTG1 UBC RPL3 RPL37 RPLP0 MT-ATP6 JUNB RPS28 RPL18 UBB MT-ATP8 RPS14 RPL39 PFN1 GAPDH HSPA1A RPL18A SRGN RPS27A RPL26 RPL19 RPS15A HLA-A DNAJB1 RPS3A CREM RPS13 MT-ND1 RPL21 RPS25 BTG2 RPL35A FAU RPL8 RPL7A RPS24 RPS6 RPS16 RACK1 NFKBIA RGS1 RPL29 CALM1 RPL9 RPL37A MT-ND5 TNFAIP3 RPS23 IL7R RPL36A PTMA NFKBIZ UBA52 EIF1 CRIP1 CORO1A RPL14 HSP90AB1 RPL10A CXCR4 RPL4 EEF1B2 RPL36 RPS9 RPL27 NACA VIM H3-3B RPS7 HSPH1 ATP5F1E HLA-E RPL17 RPSA MYL12A RPL12 CD69 TAGAP RPL35 RPS29 RPL6 SARAF ZFP36L2 MT-ND4 ARHGDIB BTG1 RPS21 EEF1D PNRC1 EEF1G HSPA5 FYB1 CD3E IFITM1 RNASEK EEF2 MT-ND2 FTL S100A4 JUN IFITM2 CYTIP OST4 LAPTM5 RPL36AL PLAAT4 PFDN5 SAMSN1 DNAJA1 EIF4A1 FXYD5 HSPE1 CTLA4 IL32 RPL24 CFL1 SAT1 RPL13A NPM1 NDUFV2 SRSF7 ITM2B POLR1D CCNI CALR ZFP36 COX7C PTPRC GADD45A CD3G RAC2 MYL6 UBE2D3 CCL20 RPS5 RPL22 TOMM7 RPL15 C12ORF57 LSP1 SON H4C3 RPL23A RPL31 DDX5 EVL AHR SERF2 STK17A CD6 PTGER4 OSTF1 OAZ1 TMA7 PHLDA1 COMMD6 PPIA TMSB10 RHOH TUBA1B KTN1 PDIA3 DUSP5 BTF3 SSR4 ATP5MC2 S100A6 CEBPD TRAV13-2 STAT1 HSPB1 PSMB9 RPL38 HNRNPA2B1 FLT3LG CSTB CDK2AP2 CSK MIF MCL1 PSME1 TXNIP RHOA HNRNPA1 UBL5 PTGES3 PDCD4 SLC2A3 ERCC1 SNHG29 NFKBID SRP14 ARPC1B NR4A2 CMTM6 CD44 ITGA4 GNB2 TRAV19 BHLHE40 HMGB1 ARPC3 TRIR PRKCH NME2 KLRB1 CLIC1 ANAPC16 NCOR1 COX5B MYL12B SAP18 RHOB TCF25 SOD1 SRSF9 FOSB PCBP1 CD96 TBC1D10C CD164 CD3D ACTR2 COX4I1 SH3BGRL3 KMT2E HSPD1 PABPC1 SPCS2 SERP1 TNFRSF1B TAGLN2 PCBP2 RPL5 LDLR YWHAB CSRNP1 COX8A MYADM TRMT112 ATP5MG GABARAP ELF1 PRR13 RPL22L1 WAKMAR2 SIT1 GADD45B RPS17 S100A10 RBBP4 HNRNPDL CD2 UQCRQ COX6B1 RGS2 CHCHD2 TAF1B RPL27A TAP1 MTDH CD99 ARF1 TBCA TMEM14B HNRNPA0 RPS11 SEC61A1 RSRP1 TUBB ACIN1 HNRNPUL1 OXNAD1 LMO4 CTSH RBM8A RGL4 NFU1 AP1G2 TSC22D3 CASS4 CCNH RPS26 PIK3R1 SH3KBP1 SSBP4 TPR GADD45G ZFAS1 COTL1 BCL7C MZT2A LAT FUS SDF2L1 KRTCAP2 MRPL41 CENPX SETD5 TMC8 COX17 DNAJB4 CDKN1B ODC1 ACAP1 ANXA6 CYBA SUMO3 TMBIM6 TMEM259 CD37 PRDX2 UBE2D2 BTG3 PARK7 CD247 TMEM258 HNRNPAB SPOCK2 SEPTIN1 LAMTOR5 RNF11 SNX5 SASH3 PPM1G ARF6 EIF5A IL4I1 RBM3 RSF1 SLA CALM3 SNRPB SHISA5 ARID5B KXD1 PTPRCAP GBP1 STN1 NSA2 ATP1A1 LCK PSMB8 PSMB3 HNRNPC RNF167 RESF1 ACTR3 CKLF CD52 RPS27L GSTK1 GIMAP7 DCXR RGS14 VASP PBXIP1 GOLGA4 RPL23 YWHAQ SNHG16 TAOK3 IGBP1 DOCK10 SRRM1 PSMC3 IKZF3 AHI1 RPL7 TMEM87A EIF5 MED1 ANKRD12 IWS1 SLC25A6 TUBA1C LRRC41 SEC11C DLST VMA21 MKNK2 RALB THUMPD1 DBF4 RHEB WASHC3 VRK2 SUPT4H1 FADD CHD1 CAPN1 TMEM179B CKS2 RGS10 EIF4E BPTF RNMT DDX39A GADD45GIP1 DYNLL1 SF3B2 CDKN2AIP RAP1B ACTN4 ABRACL MOB4 METTL23 ITSN2 MTF2 VSIR PER1 COX7A2 COX14 IPCEF1 YWHAH CYRIB MOB3A SP2 BCL11B STUB1 RNF213 PRKRA EDF1 PRPF38B RHOG BAG1 ATOX1 CACYBP MAP4K1 PPP1R2 PAXX ADA NDUFA12 WAC ELK3 PRRC2C MUS81 CCT4 ARL5B RALBP1 FBXW11 LCP1 LEPROTL1 PDCD10 HNRNPR NR3C1 PPP4C BATF ABLIM1 PSMA2 TRAT1 RUNX3 CIB1 ARPC2 ABI3 UGP2 NDUFAB1 HCLS1 DERL1 RAB5A INSIG1 CNBP CHROMR DSTYK DPP4 TOMM6 PDE4D LUC7L2 JUND ARHGDIA CEMIP2 POLE4 ANXA11 PLEKHB2 SDCBP MANF FAM83D IMP3 TMCO1 SUCLG1 MBD2 NDUFA10 PDIA6 S100A11 DNPH1 TARDBP PTDSS1 RP11-25K19 LPXN FYTTD1 RORA NASP ATP6V1F RIPK2 IDI1 RASSF7 MT-ND6 RWDD1 CYB5R3 MAZ PIM1 PSMA3 GPX4 PACS1 TNIP1 PLEKHF1 STXBP2 MSMO1 UQCR11 EPB41L4A-AS1 TMIGD2 PATL1 COX6A1 CCND3 TUBB4B OCIAD2 UBE2V2 ASMTL MAP3K8 FERMT3 SH2B1 ACP5 RPS6KB2 GABPB1-AS1 MTRNR2L12 ZFAND2A EIF3D ATP5PB NEDD8 NDUFA1 C19ORF25 CYB561D2_ENSG00000114395 PPP1R18 SYF2 PIGT SLC25A3 KHDRBS1 GIMAP4 NAXE HMGN1 RGS19 SUMO1 YWHAZ ERN1 LRRFIP1 NCF1 GALNT11 PSMG2 SNRPA1 IFNGR1 ENSA LAMTOR4 TBRG1 LINC-PINT PSMA7 SH2D2A HCST GUSB SEC61B TAX1BP1 RNF149 GNG5 PIP5K1A SIGIRR ATP6V1H RXYLT1 TRAPPC2B STK4 GPR174 CLEC2B CDC5L MLEC SQSTM1 GNAS DAXX ADAR NRIP1 TPI1 ITM2C NOSIP HADHA GTF2B ATP5IF1 ATP6V0B VEZF1 MYBL1 TOMM5 UQCRFS1 PDE4B ELL MRPL18 TRAC PIGBOS1 PPIB MDH2 PMVK PGAM1 PEX2 RPA2 FBXW5 MICOS13 CYBC1 ELOB PSMB4 COX6C GTF2I CHMP4A ALOX5AP SSR3 SELENOF AIP UXT TBCB SKAP1 RBM4 NDUFA9 DNAJB11 ATP5MK REX1BD GATA3 MVP GIMAP5 DNAJC2 SHKBP1 CAP1 SSBP1 SUZ12 SNHG6 LTB HECA WDR83OS DUT ANXA2 RRAGA PPP1R12A TAPBP DOK2 DYNLT3 RBPJ ATP6V1G1 CD40LG OCIAD1 TCEA1 ATM SLC44A2 OGA HNRNPK PPP2R5E NDUFA4 PLAUR NT5C3A RBL2 PHTF2 ZFAND5 CCR6 BCL3 GYPC SPTY2D1 MBP AKAP9 CD83 HNRNPF RAC1 SSR2 SMARCC2 MRPS21 SRSF3 DDX21 NDUFC1 BST2 FSD1 TRIB2 SF3B6 GDI2 HLTF APBB1IP LY6E NDUFB4 NAA38 FKBP1A PSMB10 FDPS MBNL1 UQCRB TMED2 H1-10 RO60 CEP57 CDKL3 CAPNS1 GIMAP1 EZR HOXB4 JAGN1 EIF2S2 DNAJB9 KCTD20 ARGLU1 EWSR1 SH3GL1 EIF3F RABGGTA MEA1 MLLT3 SARNP UBAC2 TRA2A ARL6IP1 GPI GUK1 DPH3 TGFBR2 ASB2 SRP9 GNAI2 RAB5IF PPIH EIF4A2 UBQLN1 GPR65 RASAL3 NDUFB6 RPN1 HNRNPA3 GBF1 OGG1 LCP2 DESI2 DUSP2 MGST3 CCDC18-AS1 CD53 POLM RTRAF ME2 PHAX ATP6AP1 RDH11 TCERG1 ACTR8 SLC1A5 DSE SLC25A39 FDFT1 CIAO2A LATS1 WDR33 CCDC107 MARCHF7 FAM174C LRPAP1 COG6 IL10RA MNAT1 CLIC3 TUSC2 NAT9 FAM13B EPG5 FKBP11 CREB3 ARHGAP45 MRPL9 MGA MIDEAS PHF20 CHD9 RCE1 VDAC2 EEF1E1 SSB GBP5 NMRK1 CCT8 SCAP ACER3 VAMP8 TALDO1 PHGR1 TC2N GRID2IP EIF2AK1 TOLLIP TBC1D10A C17ORF100 NFRKB CCDC137 E4F1 XRCC5 FCMR PGPEP1 DGCR6L CHURC1 OFD1 MAP4 PLEKHA2 AMD1 MORF4L1 IER5 DCTN2 CTDSP1 DDX6 AHCYL1 PNKD ELOVL1 FXR1 SAPCD1 SELENOT UBE2L6 PMM2 RNF113A PGRMC2 CPSF6 ZNF267 MRPL50 IL16 CREBL2 ASXL2 NUDCD3 ZNF445 C3AR1 CCR5 DGKE RUNX1 AKAP11 RP11-562I5 ARIH1 PPP1CB NDUFAF3 MRPS5 ACTR1B TPD52 SF3A2 CLTC TMEM33 NDUFA5 HMGXB4 PPP1R21 RC3H2 ZNF580 PRKCSH GZMM SEC13 CTNNA1 VPS18 SYNGAP1-AS1 MAPRE1 ZC3H8 DSTN MAN2B2 MGAT4A PAK2 CRY1 LSM7 ZC3H6 TMEM230 CPQ PSMD3 TIMM17A TYMP ST6GALNAC6 PHB2 KDELR2 LAMP2 NUMB RALY ADAM8 R3HDM4 DNAAF2 ADPRM CCSER2 THRAP3 KIDINS220 RBM15B DEF6 ADNP POLR1E IQGAP1 POFUT2 EIF3K NOL7 NMUR1 NOM1 GNPDA1 PKM CWC25 IGF2R ATP2B4 GPS1 LRP10 CA10 ERI3 SMAD3 ANP32B COX7A2L EHBP1 CHUK H2AZ1 LPIN2 ASAH1 MRPL33 BCL2 DNAJC1 USO1 SUMF2 CHCHD3 NAA20 TMEM60 IER3IP1 FLII YBX1 LAIR1 MESD TSC22D1 SF1 TMEM273 PLCG1 NMD3 TMEM18 ABHD14A SLFN5 ENTR1 CHMP6 IRF2BP2 CHCHD5 HTATSF1 RPS10 LSM14A SUPT7L PEX26 YJU2 HDDC3 CFDP1 SELPLG ENO1 KDM5A MTLN ZNF706 ERVK3-1 DCP1A DBNDD2 TAX1BP3 SPINK4 RNF31 NEPRO SHFL PSMD8 IL17RA CHST14 MPC1 GLOD4 BICRAL SMIM26 KIF5C PRKAR1A TENT5C NPRL2 TTC9C RPUSD3 SSH1 ARHGAP35 RSBN1 ECHDC2 G3BP2 PRR5 SMG9 SUPT3H ATP2A3 SIN3B LSM12 PTEN ZNF93 NONO SRRM2 ITFG1 C21ORF91 HINT1 GOT1 C1QTNF1 AZIN1 MZT1 LAP3 NDUFB8 NDUFA13 FRMD4B LYRM2 IL21R PARP12 SCMH1 PBRM1 HAUS4 RAB4B FLJ40194 HIKESHI NR1D1 MFSD8 SRSF1 SMIM30 XRN2 SAMD4B SEM1 CAPZB PIAS4 IKZF1 EIF3I AGTPBP1 MZT2B C1ORF122 MUTYH WDR61 DDX39B CCNDBP1 ZNF816 PFDN2 POLR2F ANO6 NDUFS6 CTNS ATP11B SIVA1 EIF4E3 CTSC SMC4 CDC42SE2 FAM131C PLPP6 TRPC4AP MTMR14 GABARAPL1 FAM133B RFNG EIF2B1 LSM1 ISY1 EIF3G ORAI2 FAM53B SEPTIN2 HTRA2 SEPTIN6 KAT7 MBD4 CAPZA2 WDR20 NDUFB7 TWF2 COMMD5 GPR68 SPPL3 TMEM59 NCL PJA2 SDHA SLMAP CBX3 TMEM161B-AS1 PPOX ACBD3 GOLGA7 GTF2H1 MAP3K2 COA6 PHF5A COPS7A NAA10 SNX14 CUL1 ALG13 TNIK STARD3 LINC00892 SNAI3 RABAC1 USP42 NDUFA3 DDX58 PSMC6 RBCK1 HNRNPH3 ZFAND6 SDHD CHCHD10 PUM1 LDLRAD4 TCHP CNOT11 KIF1C SEC61G TACC1 PSMD6 NIPBL DAPP1 RAB2A TRIM5 LMNB2 SUPT6H AP3S1 CLK1 NSUN2 NUP88 PPA1 UMAD1 SETD9 DCAF15 ANXA2R TMEM9B EGLN2 MRPS18B PSD4 GTDC1 SUCO BUD31 IER2 UPF3A PINX1_ENSG00000254093 NUDC TAF7 UTP15 PTPMT1 PSME2 MALT1 PSTPIP1 TSPYL1 CTSD HIGD2A SBF1 LINC00667 FEM1A RAB5C TAF2 SHOC2 RASL11A AMZ2 TRIM65 ATP6V0E1 HEBP1 PSMD1 NDUFAF8 CORO1B MRPL34 UBN1 RAE1 SPINT2 SOCS1 TRIM52-AS1 TFDP1 TES YIPF2 STX5 BZW1 PDCL3 NOLC1 TPM3 INTS11 RAB35 COQ10B ZNF480 C1D THAP4 EBAG9 FOXP1 SLC30A6 NUS1 MRPS24 PPP2CA USP4 C9ORF16 MSH2 UQCRC1 SLC38A1 ACO2 WASHC5 TXLNG PHF20L1 RAB6A SERINC1 YTHDC1 MIF4GD C1GALT1 KRR1 NDFIP1 GMFG DCAF8 SNX6 PAIP2 PLAA SRP19 RAD21 TECR OTUB1 CLK3 BNIP2 NEK9 COPE RAB9A RMC1 JOSD2 MAP3K1 RNF144B ADGRE5 UBA5 NLRC3 LYPLAL1 SURF6 PDE9A ZNF720 C17ORF107 SPPL2A OPTN LACTB DNAJA2 CAMK4 PLP2 ARMC6 IDH2 FUNDC2 SNHG3 CAPG TMEM106C DR1 PGGHG TUT7 CYB561D1 EIF3E FURIN VPS26B APBB1 NTAN1 ATP5MF MAX MRPL14 ATP5ME ORAI1 SEC62 TNFRSF1A HTT EIF2AK4 SCAPER POLB FAM204A RRP12 ABHD5 IMPDH1 TAF1D TDG IMP4 SMC6 TRIM28 FAS ILK PPP1R16B TSR3 NOP56 EIF5B CUTA GLO1 PPP3CB CENPK MTRR ITGAE GOLGA3 GFER EPSTI1 THEMIS TRIAP1 HSD17B7 TSTD3 ANKRD39 PIP4K2A MTFP1 FAM222B SRFBP1 HOXB2 TRAPPC5 SRSF8 HECTD1 CCAR1 XKR8 RBM26 ISCA1 TRBC2 UBE2B TIMM8B RPUSD1 SAT2 FAM102A ST6GAL1 METTL9 C2ORF68 PERP SNU13 ANKZF1 CBX5 GPN1 COMMD8 EGLN1 HERPUD1 SENP6 SLC38A2 AARS1 GNL2 G3BP1 GPR108 PSMA4 SP100 MFSD10 TMEM70 QRSL1 DAD1 RPS19BP1 SNHG1 CSDE1 TAF4B ASF1A N4BP2L2 RNF138 FBXL5 JAK1 NCLN RAB11B PPP1CA ITGB2 KDELR1 PSMB1 MTX2 SLC35E2A NDUFS8 HAGH PIGW IVNS1ABP SDHC POLDIP2 EIF4G1 ANKRD36 SEPTIN7 CSNK1D PSMF1 TM9SF4 GAB3 MYL6B HNRNPL WIPF1 GATAD1 RFFL PRKCQ-AS1 GTPBP10 INO80D IFITM3 GAS5 WASHC1 CTNNAL1 PTGER2 USP33 ARRB2 ARMCX6 GMIP UMPS TRIM22 DDX41 SFPQ RBM25 MRPL38 CNP FAM193B MAP3K11 WDR74 PIGC CDC37 PPIP5K2 TRMO AGTRAP JMJD8 ARID5A ALKBH7 CYTH1 SET PAF1 SNHG5 FAM174B ADI1 SUDS3 CDCA7 GTSF1 MTCH1 TGIF2 MLX PYGO2 EIF2S1 DNPEP IL23R RPA1 CDC37L1 GYG1 ANP32A SRSF2 RNASEH2B WDR44 ZPR1 GIPC1 RRM2B ILVBL ARL2BP BRK1 YIF1A CEP164 DPY30 NUDT7 PPP2R1A R3HDM1 MSN PSIP1 HMGN4 CCDC152 STOML2 EEF1AKMT2 H2AZ2 CHST12 CHFR FKBP15 SPNS1 APH1A NR4A3 CNOT9 GPSM3 TESK1 CDK5RAP1 TNFAIP1 FAM118B ZBTB7A PCED1B-AS1 MRPS18A BUB3 ICMT DDIT3 DNM1L DCTN4 PPIG CCDC91 LZIC CHD2 BIRC2 AHSA1 ARAP2 COMMD10 NUCKS1 ASAP1 DCTPP1 ETFDH NAA80 STAU2 CCNL1 CASP3 CCDC174 PRNP APOA1-AS PPP1R11 SRP68 C1ORF43 RASSF1 MYD88 AKAP13 SLC25A5 ERH UBE2I ATP5PD ORC3 NUP37 TTC31 WBP2 PDE7A KPNB1 ETFA STK25 PJVK COX19 ATP5MC3 PPP2R2A MDH1 PSMB6 GTF2E2 B3GNT2 TSC22D4 GPD2 MBIP ZFR JTB GPRIN3 LPIN1 SLC35A1 NXT1 KLHL22 DYNC1I2 AP1S1 PPP1R15B SMIM15 TMEM65 HAUS3 KLHL28 MRPL20 MVK ATR MIB2 YPEL3 SCOC WASHC2C SPRY1 LCLAT1 CCDC25 SLC43A3 CIAO3 WBP11 TRIM4 C17ORF49 DTX1 VAMP5 ZSCAN16-AS1 TNFSF13B ARHGAP15 CISH DCP2 GRPEL2 PHB AASDHPPT PUM2 ISG20 PFKM SLC2A1 NBL1 TMED4 COPS3 SRSF11 ZNF710 POMP TTC14 ZNF414 PISD TMEM219 MIR23AHG COQ4 OXA1L DNAJA3 KMT2A MAGOH LYPLA2 WSB1 POLR2B AKIRIN2 PSMB2 UIMC1 DDX18 TCOF1 TNFRSF14 STARD3NL CELF2 NABP1 PTP4A2 LRRC59 TMEM167B CAMK2G SLC4A1AP EIPR1 NUCB2 MED6 VEZT SP4 ARL3 USP12 ARHGAP9 VMAC FIP1L1 RP11-316M20 ADH5 TMC6 IRF1-AS1 LIG4 SNX1 MACF1 GNL1 TPM4 RMND5A LPAR6 DNAJB2 NOP53 CDK2AP1 FIS1 PITPNM1 SF3B4 SNAP47 S100PBP C4ORF33 CHST11 RGS18 C7ORF31 PSMD4 EFCAB7 TTC39C ANXA5 KRI1 FAM214B DUSP12 NR2C1 CBLB UNG LSM8 MYO9B GLMN RALA GTF3C6 ISCA2 DNTTIP2 SELENOK MRPS36 KPNA3 THRA PTS DUSP6 VTA1 PPARG SEPTIN9 UBE2N FRY BCAS2 CAMKMT SAMHD1 GBP2 TMEM50A STAT5A IL17A PCYT2 ISG15 MAPK3 PDCD4-AS1 TIAL1 RTF1 CSNK2B RBM17 LPCAT4 ST13 CNBD2 TMEM184B RNPC3 ARCN1 SMAP2 BLOC1S1 NAMPT ZNF597 EMG1 SLC66A3 HDLBP C11ORF58 SH3BP1 SMARCB1 RP11-726G1 ZNF564 IFI27L2 ATG3 C6ORF226 REL GIMAP6 OLFM2 OGFR BAZ1A RAB7A EIF3L TBC1D20 SDHAF2 PARP10 GAK FYN TBXAS1 PRDM1 FAM241A TMEM263 SRPK1 CD2BP2 SERGEF SURF4 MTERF4 PPP1R10 IRF9 USP14 ZFAND2B AK3 CYTH3 PEBP1 PLIN2 NARF ICAM2 ZRSR2 HELZ C14ORF119 DCAF5 RBM42 ITFG2-AS1_ENSG00000258325 ZNF524 CRNKL1 FASTK UBXN1 MYO1G QTRT1 FLOT1 DNASE2 STK40 USP47 RFLNB ZFYVE28 FHL1 TATDN2 SLC25A25 ATP6V1B2 GSAP MRPL10 PSMD11 ZNF821 RANBP9 TIPRL TMED10 EIF3M HMGN2 UHMK1 TRBV30 EMP3 HPS6 TTC17 OSBPL8 COMT ERP29 GRPEL1 ZNF595 ROGDI RAD23A LYSMD2 ANKRD44 FRMD8 SRSF10 BCL2A1 TBC1D31 ALDH9A1 CHMP1A GABARAPL2 RPS20 MRPS2 CAST TMEM200A H4C2 ATRAID CRYBG1 CRY2 OSBPL7 TRAPPC1 POLR2K TESMIN UBE2D1 GNAI3 SLC39A6 HMGCS1 TMEM138 TMEM223 CHP1 ATP5MJ MSRA SQLE ATRX ZNF207 PPP1R15A DNAJC3 SNX3 FBXO9 FAM162A BTF3L4 NKIRAS1 PRDX6 EHMT1 SYS1 USP36 STX4 TMED7 NR4A1 GNA13 LINC00299 ST3GAL5 IDH3A NAA16 SLC38A10 CRYBB2 B4GALT1 METTL26 MRPL54 TRAF3IP2 DTX3L NDUFA6 NDUFAF4 SUCLA2 CALHM2 GTF2IRD2B RP11-706O15 IK MAGI3 TCTA PBDC1 BRD7 IRF1 P2RY10 CTBS CUX1 PRKD2 PGK1 BCLAF1 IRAG2 CIAPIN1 TNFRSF25 SNRNP200 IL2RG PSMA6 UQCRC2 RTF2 SLC35C2 GNG2 ZBTB22 MEF2A NOP10 LY9 MAP2K2 COMMD1 TMEM203 VPS29 SOD2 CITED2 FLI1 FMNL1 MRTO4 RETREG2 FGFR1OP2 RNF41 UTP4 TKT DHX36 ID2 TMEM245 ERP44 ESF1 BAG3 NDEL1 DAP3 HM13 PTGR2 ZFAND3 UCP2 BIN2 MCF2L2 PHKB SUB1 PDE3B CEP250-AS1 RNH1 WDR1 FAM117A CAMTA2 UGCG LINC00294 PHF14 PTPN22 KATNBL1 NT5DC1 TRIM38 ZNF780A C19ORF53 SNRPG TDP2 EIF4H TMEM140 ZCCHC17 ANKLE2 CDS2 EPHB6 CCDC90B ATAD3C COL18A1 FBP1 AHCY AFTPH ZNF432 HSPA9 CMSS1 FGD3 ANXA7 SLC25A13 RNF5 KRT10 KDM2B-DT ATP5F1D SCAMP3 KIF2A SHMT1 MRPL57 CD5 MLH1 ACOX3 ZNF593 TTI2 TRAPPC10 SRP72 CXCR6 COPS5 TET2 HEBP2 OGT ANKRD13D CEBPZ MRPL23 UBE2V1 AP2M1 SLC39A4 LLGL1 DLD ACAA2 SCAMP4 PARP8 CXXC1 MAD2L2 GIT2 OSGIN2 FAM219A ERLEC1 PTPN7 TMEM41B SERTAD3 TCF7 ARPC5 NCBP2AS2 ZNF211 SNRPB2 B9D2 ANKRD27 TIPARP RAB1B SRM PARP16 EIF4EBP2 GDI1 ARPC5L CEACAM21 IVD EXOSC4 KDM6A STARD7 JAK3 CBR1 ARL6IP5 STK26 UFM1 ARHGAP18 ESD ISOC2 HCFC1 RING1 PRELID1 MGME1 NDUFC2 CELF1 PNP IFT27 KBTBD2 TUFM TSTD1 CISD3 C17ORF67 BUD23 SCRN2 CHCHD7 FARSA RP11-29H23-1 LINC02481 TPP1 THYN1 ARHGEF1 GCLM PRPF39 TRAPPC6B SLF1 SMIM3 EVI2A AZIN2 ASH1L API5 COG2 SLC25A22 CHD4 WRAP73 CDC26 UBE2L3 EPC1 ATP2B1 CCAR2 RRN3 UBA1 CSNK1A1 CASP8 PELI1 CNOT2 SRP54 VAPA MRPS33 NDUFB11 NAF1 UBE2G2 DVL3 TLE5 CALCOCO2 UQCR10 SIRPG COX7B NDRG1 EIF2S3 H3C4 SMCHD1 RAN MAPKAPK5-AS1 TSG101 HAX1 ANKDD1A DCTN3 CD55 PSKH1 PET100 IMMT SRGAP2C H3-3A PIGH RBM48 FAM177A1 TEFM FBXL4 MID1IP1 PRR5L SLC25A11 NGDN NCAPD3 SLC9A3R1 CDIP1 GRK6 SDF4 DENND2D CASP7 EFR3A NFAT5 PRELID3B TEX30 CIRBP CSNK1G1 MCRS1 CHMP2B TMOD3 PRKX POLR2C NSFL1C POLR1H TAP2 PIANP ARL6IP6 RIPK1 TDP1 SCAF11 SNRPN ZNF335 USP3 UBQLN2 ARMCX3 HSDL1 RSU1 TXNL1 TMEM101 BAD KIF5B AP1AR FAM104A GBP4 AC058791 NDUFB1 DPP7 NDUFB2 RANGAP1 SETX MAF1 NME4 TRGV10 TACC3 AP5Z1 PITPNB EDEM1 CTD-2095E4 ANKRD28 MINDY2 YPEL4 IGHV4-34 TRIM44 ARL6IP4 GART RAD18 SF3B5 BZW2 ECD ZYX HECW2-AS1 CNTRL RECQL PARP9 UPP1 PYCR2 ARHGAP30 TBC1D22B RASGRP1 STIM1 ZNF263 ZNF230 LDLRAP1 PDCD2 FOXO1 UNC119B RNF114 YPEL5 IFNGR2 PRRC1 EIF4B PDXK RAB8B RICTOR ARL8A ESCO1 RBM38 PPIF GMPS HDDC2 FBXO42 PTP4A3 SERPINH1 TRIM69 WTAP PEMT VPS72 WAS PPP4R3A CDV3 CETN2 VAC14 TRIM21 DAZAP2 ARF5',\n",
       " 'cell_type': 'CD4-positive helper T cell',\n",
       " 'tissue': 'ileum',\n",
       " 'batch_condition': 'A29',\n",
       " 'organism': 'Homo sapiens',\n",
       " 'sex': 'female'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "arrow_ds[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61148588",
   "metadata": {},
   "source": [
    "# Group cells into multi-cell samples\n",
    "\n",
    "Here, we group together cells with the same tissue label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d2d9b29a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_multi_cell_groupings(hf_ds):\n",
    "    batch_tissue_to_cell_indices_dict = {}\n",
    "    for sample_idx in range(hf_ds.num_rows):\n",
    "        # Load sample, get batch sample and tissue type\n",
    "        sample = hf_ds[sample_idx]\n",
    "        batch_sample = sample[\"batch_condition\"]\n",
    "        tissue_type = sample[\"tissue\"]\n",
    "\n",
    "        # If new batch sample and tissue type combination is found, add to dictionary\n",
    "        if (batch_sample, tissue_type) not in batch_tissue_to_cell_indices_dict:\n",
    "            batch_tissue_to_cell_indices_dict[(batch_sample, tissue_type)] = []\n",
    "\n",
    "        # Add sample index to dictionary\n",
    "        batch_tissue_to_cell_indices_dict[(batch_sample, tissue_type)].append(sample_idx)\n",
    "    return batch_tissue_to_cell_indices_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "89acb513",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_tissue_to_cell_indices_dict = get_multi_cell_groupings(arrow_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a5edfe42",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(batch sample, tissue type): [list of cell indices with this combination]\n",
      "('A29', 'ileum'): [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]\n",
      "('A29', 'lung'): [346, 347, 348, 349, 350, 351, 352, 353, 354, 355]\n",
      "('A29', 'thoracic lymph node'): [3547, 3548, 3549, 3550, 3551, 3552, 3553, 3554, 3555, 3556]\n",
      "('A29', 'mesenteric lymph node'): [8117, 8118, 8119, 8120, 8121, 8122, 8123, 8124, 8125, 8126]\n",
      "('A29', 'bone marrow'): [10955, 10956, 10957, 10958, 10959, 10960, 10961, 10962, 10963, 10964]\n",
      "('A29', 'skeletal muscle tissue'): [12365, 12366, 12367, 12368, 12369, 12370, 12371, 12372, 12373, 12374]\n",
      "('A29', 'liver'): [12381, 12382, 12383, 12384, 12385, 12386, 12387, 12388, 12389, 12390]\n",
      "('A29', 'spleen'): [13041, 13042, 13043, 13044, 13045, 13046, 13047, 13048, 13049, 13050]\n",
      "('A31', 'spleen'): [16664, 16665, 16666, 16667, 16668, 16669, 16670, 16671, 16672, 16673]\n",
      "('A31', 'ileum'): [18835, 18836, 18837, 18838, 18839, 18840, 18841, 18842, 18843, 18844]\n",
      "('A31', 'omentum'): [18868, 18869, 18870, 18871, 18872, 18873, 18874, 18875, 18876, 18877]\n",
      "('A31', 'lung'): [18898, 18899, 18900, 18901, 18902, 18903, 18904, 18905, 18906, 18907]\n",
      "('A31', 'liver'): [20155, 20156, 20157, 20158, 20159, 20160, 20161, 20162, 20163, 20164]\n",
      "('A31', 'thoracic lymph node'): [22211, 22212, 22213, 22214, 22215, 22216, 22217, 22218, 22219, 22220]\n",
      "('A31', 'caecum'): [25214, 25215, 25216, 25217, 25218, 25219, 25220, 25221, 25222, 25223]\n",
      "('A31', 'bone marrow'): [25285, 25286, 25287, 25288, 25289, 25290, 25291, 25292, 25293, 25294]\n",
      "('A31', 'thymus'): [25648, 25649, 25650, 25651, 25652, 25653, 25654, 25655, 25656, 25657]\n",
      "('A31', 'mesenteric lymph node'): [26002, 26003, 26004, 26005, 26006, 26007, 26008, 26009, 26010, 26011]\n",
      "('A31', 'skeletal muscle tissue'): [28373, 28374]\n",
      "('A31', 'duodenum'): [28375, 28376, 28377, 28378, 28379, 28380, 28381, 28382, 28383, 28384]\n",
      "('A29', 'caecum'): [29110, 29111, 29112, 29113, 29114, 29115, 29116, 29117, 29118, 29119]\n",
      "('A29', 'transverse colon'): [29378, 29379, 29380, 29381, 29382, 29383, 29384, 29385, 29386, 29387]\n",
      "('A29', 'sigmoid colon'): [29640, 29641, 29642, 29643, 29644, 29645, 29646, 29647, 29648, 29649]\n"
     ]
    }
   ],
   "source": [
    "print(\"(batch sample, tissue type): [list of cell indices with this combination]\")\n",
    "for key in batch_tissue_to_cell_indices_dict.keys():\n",
    "    print(f\"{key}:\", batch_tissue_to_cell_indices_dict[key][:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "208c2009",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_multi_cell_arrow_ds(hf_ds, batch_tissue_to_cell_indices_dict):\n",
    "    \"\"\"\n",
    "    Function to take in a Huggingface cell sentence arrow dataset and a dictionary \n",
    "    mapping a tuple of (batch sample, tissue type) to a list of cell indices, and \n",
    "    return a new arrow dataset containing the multi-cell groupings.\n",
    "\n",
    "    Arguments:\n",
    "        hf_ds: Huggingface arrow dataset containing cell sentences to group into multi-cells.\n",
    "        batch_tissue_to_cell_indices_dict: Dictionary mapping a tuple of (batch sample, tissue type) to a list of cell indices.\n",
    "    \n",
    "    Returns:\n",
    "        multi_cell_groupings_ds: Huggingface arrow dataset containing multi-cell groupings.\n",
    "    \"\"\"\n",
    "    # Initialize list to store multi-cell groupings, each element is a list of 5 cell indices\n",
    "    multi_cell_groupings_list = []\n",
    "    batch_sample_labels = []\n",
    "    tissue_type_labels = []\n",
    "    organism_labels = []\n",
    "\n",
    "    # Loop through each sample in original dataset - keeps number of overall samples (roughly) the same\n",
    "    for sample_idx in tqdm(range(hf_ds.num_rows)):\n",
    "        sample = hf_ds[sample_idx]\n",
    "        batch_sample = sample[\"batch_condition\"]\n",
    "        tissue_type = sample[\"tissue\"]\n",
    "\n",
    "        # Retrieve list of cell indices for this batch sample and tissue type\n",
    "        sample_key = (batch_sample, tissue_type)\n",
    "        all_cell_indices = batch_tissue_to_cell_indices_dict[sample_key]\n",
    "        if len(all_cell_indices) <= 5:\n",
    "            # If there are less than 5 cells for this batch sample and tissue type, we will sample with replacement.\n",
    "            #  Note that another option here would be to just skip this sample.\n",
    "            sampled_cell_indices = random.choices(all_cell_indices, k=5)\n",
    "        else:\n",
    "            # If there are 5 or more cells for this batch sample and tissue type, we will sample without replacement\n",
    "            sampled_cell_indices = random.sample(all_cell_indices, k=5)\n",
    "\n",
    "        # Add list of sampled cell indices to multi-cell groupings list\n",
    "        multi_cell_groupings_list.append(sampled_cell_indices)\n",
    "        batch_sample_labels.append(sample[\"batch_condition\"])\n",
    "        tissue_type_labels.append(sample[\"tissue\"])\n",
    "        organism_labels.append(sample[\"organism\"])  # all cells should be of the same organism\n",
    "\n",
    "    multi_cell_groupings_ds = Dataset.from_dict({\n",
    "        \"multi_cell_groupings\": multi_cell_groupings_list,\n",
    "        \"batch_label\": batch_sample_labels,\n",
    "        \"tissue\": tissue_type_labels,\n",
    "        \"organism\": organism_labels,\n",
    "    })\n",
    "    return multi_cell_groupings_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "61064af0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 29773/29773 [00:02<00:00, 12727.38it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['multi_cell_groupings', 'batch_label', 'tissue', 'organism'],\n",
       "    num_rows: 29773\n",
       "})"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "multi_cell_arrow_ds = get_multi_cell_arrow_ds(arrow_ds, batch_tissue_to_cell_indices_dict)\n",
    "multi_cell_arrow_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "34cf2803",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'multi_cell_groupings': [225, 59, 3, 46, 298],\n",
       " 'batch_label': 'A29',\n",
       " 'tissue': 'ileum',\n",
       " 'organism': 'Homo sapiens'}"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "multi_cell_arrow_ds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1c844580",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cell 225: A29 ileum\n",
      "Cell 59: A29 ileum\n",
      "Cell 3: A29 ileum\n",
      "Cell 46: A29 ileum\n",
      "Cell 298: A29 ileum\n"
     ]
    }
   ],
   "source": [
    "for cell_idx in multi_cell_arrow_ds[0][\"multi_cell_groupings\"]:\n",
    "    print(f\"Cell {cell_idx}:\", arrow_ds[cell_idx][\"batch_condition\"], arrow_ds[cell_idx][\"tissue\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8428d5c9",
   "metadata": {},
   "source": [
    "# Prompt Formatting for Natural Language Interpretation\n",
    "\n",
    "We will format prompts that provide the model with multiple cell sentences and ask it to generate a summary of the biological information contained within them. This is a powerful feature of C2S, allowing researchers to gain high-level insights from their data in a human-readable format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b3c8d896",
   "metadata": {},
   "outputs": [],
   "source": [
    "from cell2sentence.prompt_formatter import C2SMultiCellPromptFormatter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "16092cfc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cell2sentence.prompt_formatter.C2SMultiCellPromptFormatter at 0x1550aaab0610>"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "task_name = \"natural_language_interpretation\"\n",
    "prompt_formatter = C2SMultiCellPromptFormatter(\n",
    "    task=task_name,\n",
    "    top_k_genes=200\n",
    ")\n",
    "prompt_formatter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "cda801f0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'multi_cell_groupings': [225, 59, 3, 46, 298],\n",
       " 'batch_label': 'A29',\n",
       " 'tissue': 'ileum',\n",
       " 'organism': 'Homo sapiens',\n",
       " 'abstract': ''}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# We would like to generate natural language summaries at inference for groups of cells.\n",
    "# For the input sample, we will put an empty string as a placeholder for the abstract, so that\n",
    "#  prompt formatting code is not affected.\n",
    "multi_cell_arrow_ds = multi_cell_arrow_ds.add_column(\"abstract\", [\"\"] * multi_cell_arrow_ds.num_rows)\n",
    "multi_cell_arrow_ds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f5e51eef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['sample_type', 'model_input', 'response'],\n",
       "    num_rows: 29773\n",
       "})"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "formatted_ds = prompt_formatter.format_hf_ds(\n",
    "    hf_ds=arrow_ds,\n",
    "    multi_cell_indices_ds=multi_cell_arrow_ds\n",
    ")\n",
    "formatted_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "8ff9f0b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'sample_type': 'natural_language_interpretation',\n",
       " 'model_input': 'Consider these 200 highly expressed genes, ordered by descending expression, from 5 Homo sapiens cells. Using this, your task is to generate an abstract summary which summarizes the biological insights contained in these cells.\\nCell 1:\\nTMSB4X MT-CO1 CCL4 B2M MT-CO2 MALAT1 ACTB RPLP1 MT-CYB MT-CO3 HSP90AA1 MTRNR2L12 MT-ND3 EEF1A1 CCL5 MT-ATP8 HSPA1B RPS12 MT-ATP6 RPL41 RPL10 RPL28 HLA-B MT-ND4L RPS19 RPS27 RPL13 HLA-C TPT1 SH3BGRL3 JUN MT-ND2 MT-ND1 RPL37 RPS18 RPL19 RPS4X RPL30 RPS2 RPS15A IL32 HSPA8 ACTG1 CD69 XCL1 RPL26 DNAJB1 RPS15 KLF6 RPS28 RPL39 TMSB10 S100A4 RPL34 RPS14 RPL29 FTH1 GAPDH RPLP2 BTG1 MT-ND4 ATF3 RPL18A RPS3 FAU RPS27A RPS3A RPL32 RPL12 RPL21 NKG7 RPL37A RPS21 RPS24 PFN1 RPS13 RPL3 RPL11 IFITM1 CD8A RPL9 RPSA RPS23 CD3D MT-ND5 TNFAIP3 RPL35A RPS6 RPL18 RPLP0 RPL27 ZFP36 IFNG PTMA RPL8 RPS25 TRBV4-1 HLA-A CD52 PHLDA1 RPL17 VIM RPL14 RPS8 UBB H3-3B TSC22D3 SRGN RPL23A RPL35 RPS7 HSPA1A RPS16 RGS2 RPL36 CORO1A RPL13A GZMA RACK1 RPL15 RPL4 NR4A3 RAC2 RPS9 RPL7A ARHGDIB CD3E UBC ANXA1 EIF1 HLA-E RPS29 CLIC1 FTL SERF2 CALM1 SPRY1 RPL6 ZFP36L2 INSIG1 ZFP36L1 CD8B UBA52 PPIA STAT1 PLAAT4 CD160 LAPTM5 DUSP1 RPL24 ATP5MC2 NR4A1 COTL1 ARL4C CD7 EEF2 HNRNPA2B1 ATP5F1E ARL6IP5 SLC25A5 MYL12B HNRNPA1 HCST MYH9 MSN EEF1G BTG2 TRAPPC1 NACA HMGN2 GNG2 MYL12A RGL4 LMO4 CRIP1 PNRC1 PPIB EEF1D SUMO2 NEAT1 FOS SARAF ETS1 TNF SNRPD2 RPS5 STK17B RPL36A ARPC2 RHOB SEPTIN1 SAT1 CYTIP EVL GUK1 MIF JUND CFL1 COX6B1 C12ORF57.\\nCell 2:\\nRPLP1 CCL4 TMSB4X B2M RPL41 ACTB MT-CO1 MALAT1 RPL13 RPS19 RPL28 EEF1A1 RPS12 CCL5 HLA-B MT-CO2 TPT1 RPS27 MT-CO3 RPL32 MT-ND3 HSPA1A XCL1 RPS4X KLF6 DNAJB1 RPS27A RPS3 RPL10 RPS15A RPLP0 RPS18 RPS28 RPS23 MT-ATP8 RPL11 RPL30 FOS RPS2 RPS7 NR4A2 RPS14 RPS6 HLA-C RPL39 RPL18A RPL37 RPL19 JUN SH3BGRL3 TMSB10 MT-CYB RPL34 PFN1 RPL26 RPL15 HSP90AA1 TRBV9 RPS15 RPS25 RPL29 RPL17 RPS24 VIM RPS29 MT-ND4L ACTG1 RPL7A RPL36 RPL18 RPS21 RPS3A RPS8 IFITM1 H3-3B HLA-A PPP1R15A PTMA MT-ATP6 FTH1 BTG1 GAPDH RPL23A RPL3 ATF3 CD74 DUSP1 UBA52 RPL35 PNRC1 RPS16 RPL37A ATP5F1E RACK1 RPL21 RPLP2 RPL9 RPL12 RPL10A HSPA8 RPS13 EIF1 RPL8 MT-ND1 SERF2 RPL38 TNFAIP3 ARHGDIB SRGN RPL36A RPSA RPS5 IFITM2 SPRY1 RPL6 PPIA CFL1 FAU CD7 EEF2 RPL35A HCST MT-ND5 SLA IL32 CORO1A NACA RPL24 HSPA1B CRIP1 HNRNPA2B1 H3-3A CYTIP RPL27 ZFP36L2 RPL22 CD52 MT-ND4 TRAV19 ZFP36L1 CKLF RPS9 PPIB RPL14 EVL COX6B1 MYL6 RPL4 ZFP36 UBC RHOB FTL IL7R EEF1B2 UBB RPL27A TOMM7 GYPC EEF1G RAC2 CAPZB OAZ1 MAP3K8 PSMB10 HSPE1 RPL5 CD3E SARAF CD99 PABPC1 SRP14 RBM3 CHCHD2 CD3D CDKN1A NDUFS5 MYL12B EIF3F TAGAP RHOH PFDN5 SRSF7 RPL31 NR4A1 MCL1 MYL12A OST4 HMGB1 RAB9A S100A10 SAMSN1 NKG7 FYB1 S100A4 NFKBIA CNBP MT-ND2 PLAAT4 DYNLL1 SLC2A3.\\nCell 3:\\nCCL4 MALAT1 CCL4L2 B2M RPLP1 MT-CO2 EEF1A1 MT-CO1 CCL5 JUN TMSB4X IFNG FOS MT-CO3 RPL41 MT-ND4L DNAJB1 HSP90AA1 RPS12 MT-CYB ACTB RPL28 KLF6 HSPA1A PTMA TPT1 RPS19 RPS4X MT-ND3 CD69 RPL30 RPS27 RPS27A RPS28 RPL13 HLA-B RPL10 RPL32 RPS15A MT-ATP8 RPLP2 RPS3 RPL19 RPS18 RPL37 RPL18A RPS21 RPS3A RPL18 RPS14 PPP1R15A ANXA1 MT-ND1 H3-3B RPS15 RPL34 RPL39 SH3BGRL3 MT-ATP6 NR4A1 RPS8 RPL35A MT-ND2 RPS25 RPS23 ATF3 DUSP1 GAPDH RPL7A RPL26 RPL8 RPS13 RPL29 MT-ND4 RPS7 TMSB10 HLA-E IL32 ACTG1 RPL12 TRBV7-9 HSPA1B RPL14 CD8A UBB RPL3 RPS29 RPS24 RPS16 RPL36A HLA-C RPL5 HSPA8 S100A4 RPL22 RPS2 TMA7 NKG7 TNFAIP3 RPL36 RPL21 RPS5 MT-ND5 RPLP0 ATP5F1E RPL11 HSP90AB1 FAU TNF CKLF RPS6 JUND RPL9 FOSB FYB1 MYL6 FTH1 CD3E BTG1 RPL6 MYO1E ZFP36 PTPRCAP KLRB1 MCL1 VIM HCST RPL37A CD7 RPL4 FTL COX4I1 SRGN EIF1 RPS9 IFITM2 PCBP2 ITM2B NFKBIZ EEF1D SARAF RPL13A CDKN1A NFKBIA PNRC1 CFL1 MYL12A OAZ1 UBA52 IFITM3 TAGAP LY6E RPL24 ARHGDIB LDLRAD4 RPL35 TSC22D3 GZMA CRIP1 JUNB CD3D PFN1 UQCR11 NEU1 NACA RPL27 RPL15 BTF3 PCBP1 ATP1A1 RPL23 HLA-A COX7A2 HNRNPA3 MRPL18 ARGLU1 PABPN1 IFITM1 ZFP36L1 MYADM CORO1A MIF IL2RG SNRPB2 RACK1 HNRNPK CHCHD2 TPM3 ANXA11 HSPA5 SLC25A6 EEF2 CSK CALR RPSA ELF1 RBM39 STK17A COMMD6 CD2.\\nCell 4:\\nB2M MT-CO1 TMSB4X MT-CO2 RPLP1 HSP90AA1 MALAT1 HSPA8 MT-CO3 EEF1A1 H3-3B HSPA1A RPL41 MT-CYB ACTB RPS27 HSPA1B RPS18 KLF6 RPL28 RPS19 HLA-B MT-ATP6 RPL10 RPL13 MT-ATP8 RPS12 RPS28 RPL11 TPT1 MT-ND3 RPS27A CCL5 JUNB RPL39 DNAJB1 RPS2 UBC ANXA1 RPS21 MT-ND4 RPL32 MT-ND4L RPS15 RPS3 ATF3 RPL34 DUSP1 RPL21 TNFAIP3 RPLP2 RPS3A RPL37 RPL30 RPS15A BTG1 PTMA MT-ND1 RPS23 PFN1 NR4A2 RPS4X RPS14 TRBV4-1 HLA-C RPS6 HLA-A RPL12 RPL3 RPL29 PNRC1 SARAF RPL19 RPL18A RPL36 VIM RPS8 RPS7 NR4A1 TMSB10 RPSA RPL26 MT-ND5 RPS29 RPS13 RPL15 EIF1 RPL17 RPL35A CSF1 FAU GAPDH HSP90AB1 MYL12A S100A4 RPL14 FTL CD69 ACTG1 RPL18 SH3BGRL3 RPL7A CD8A ADGRE5 GZMA JUN PLAAT4 PTPRCAP NEAT1 RPLP0 RPS25 RPL22 DNAJA1 CD7 RPS24 SELENOK RPL24 RPS9 RPL5 RPS5 FOS LGALS1 MT-ND2 RPL37A STK17B NACA EEF1B2 SAT1 CYTIP RPL23A JUND CORO1A RPL38 RPL13A RPL8 PPIA FYB1 CFL1 SRGN ZFP36L2 RHOH EEF1G RPS16 HSPE1 EEF2 IL32 KLRB1 RPL9 HSPH1 TSC22D3 UBA52 SEPTIN6 PFDN5 RPL10A CD3G YBX1 CD3E HLA-E RPL4 ARHGDIB PTGER4 LIMD2 CACYBP HSPB1 SERF2 FOSB HNRNPA0 ITM2B CKLF FERMT3 RACK1 MCL1 CRIP1 RPL35 PTGES3 CAPZB RPS26 SDCBP TPI1 RAC1 CAP1 COX7C RPL36A FRMD4B TOMM7 PPP1R15A IFITM1 HSPD1 GNB2 IL7R PMAIP1 ARPC3 ATP5F1E CD3D MTRNR2L12 LSP1 MYL6 CYCS TIPARP SPATS2L.\\nCell 5:\\nMT-CO2 GNLY MT-CO3 MT-ND3 MT-CYB MT-CO1 RPLP1 MALAT1 TMSB4X MT-ATP6 EEF1A1 RPL41 MT-ATP8 HSP90AA1 RPS19 MT-ND1 RPS27 MT-ND4L GZMA CCL5 RPL13 MT-ND2 RPL30 RPL28 RPS23 RPS18 RPS27A RPL10 RPS15A ACTB B2M MT-ND4 RPS4X RPS12 RPS3 HSPA1A PTMA TPT1 RPL32 RPS8 RPS14 RPS6 RPL26 RPL34 RPL39 RPS28 GAPDH CD7 RPS24 RPL18A RPS3A RPL11 RPL8 RPS15 RPL7A RPL3 RPS7 KLF6 RPL35A RPL12 RPL23A RPL19 NKG7 RPL6 RPLP2 RPS9 RPLP0 CD52 RPL29 RPL37 TMSB10 RPL37A RPS13 RPL35 DNAJB1 MT-ND5 RPS25 RPS2 RPL15 FTL TYROBP RPS21 EEF1B2 HSPA1B RPL36A HSPA6 UBC ACTG1 MTRNR2L12 ZFP36L1 RPL18 FAU RPL22 RPL9 RPS29 NACA FTH1 NR4A1 RPL21 RPL17 ATP5F1E RPS5 RACK1 PABPC1 RPL10A IER5 CORO1A RPL14 RPL36 PFN1 H3-3B SRGN TAGAP MYL12A IFITM2 RPS16 RPL27 HCST VIM HLA-B HOPX CFL1 ATF3 SERF2 SH3BGRL3 NPM1 HLA-A EEF1D NR4A3 IFNG RPL27A CKLF EEF1G UBB RPL4 H2AJ HSP90AB1 EIF1 YBX1 PLAAT4 TRDC RPS26 BTG1 GEM CD3E COTL1 HSPA8 RPL24 FCER1G PNRC1 ARHGDIB RPL5 TOMM7 IFITM1 RPL38 RAC2 CYBA TRA2B MIF SEC11A CD247 RPL7 RPL31 OAZ1 PFDN5 CD69 KRT86 RPL13A RPSA EEF2 BTF3 COX7C MYL6 RHOC UBXN4 UBA52 DBI HMGN2 GSTK1 HSPB1 KLRB1 KIR2DL4 HLA-E FOSB CIRBP PTPN22 PHLDA1 TMA7 SYTL3 ATP5MG RHOA GNAS GYPC HLA-C H3-3A COX6C KLRC2 PARK7 NFKBIA ATP5MC2.\\n.\\nAbstract summary:',\n",
       " 'response': '.'}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "formatted_ds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "285ca7f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Consider these 200 highly expressed genes, ordered by descending expression, from 5 Homo sapiens cells. Using this, your task is to generate an abstract summary which summarizes the biological insights contained in these cells.\n",
      "Cell 1:\n",
      "TMSB4X MT-CO1 CCL4 B2M MT-CO2 MALAT1 ACTB RPLP1 MT-CYB MT-CO3 HSP90AA1 MTRNR2L12 MT-ND3 EEF1A1 CCL5 MT-ATP8 HSPA1B RPS12 MT-ATP6 RPL41 RPL10 RPL28 HLA-B MT-ND4L RPS19 RPS27 RPL13 HLA-C TPT1 SH3BGRL3 JUN MT-ND2 MT-ND1 RPL37 RPS18 RPL19 RPS4X RPL30 RPS2 RPS15A IL32 HSPA8 ACTG1 CD69 XCL1 RPL26 DNAJB1 RPS15 KLF6 RPS28 RPL39 TMSB10 S100A4 RPL34 RPS14 RPL29 FTH1 GAPDH RPLP2 BTG1 MT-ND4 ATF3 RPL18A RPS3 FAU RPS27A RPS3A RPL32 RPL12 RPL21 NKG7 RPL37A RPS21 RPS24 PFN1 RPS13 RPL3 RPL11 IFITM1 CD8A RPL9 RPSA RPS23 CD3D MT-ND5 TNFAIP3 RPL35A RPS6 RPL18 RPLP0 RPL27 ZFP36 IFNG PTMA RPL8 RPS25 TRBV4-1 HLA-A CD52 PHLDA1 RPL17 VIM RPL14 RPS8 UBB H3-3B TSC22D3 SRGN RPL23A RPL35 RPS7 HSPA1A RPS16 RGS2 RPL36 CORO1A RPL13A GZMA RACK1 RPL15 RPL4 NR4A3 RAC2 RPS9 RPL7A ARHGDIB CD3E UBC ANXA1 EIF1 HLA-E RPS29 CLIC1 FTL SERF2 CALM1 SPRY1 RPL6 ZFP36L2 INSIG1 ZFP36L1 CD8B UBA52 PPIA STAT1 PLAAT4 CD160 LAPTM5 DUSP1 RPL24 ATP5MC2 NR4A1 COTL1 ARL4C CD7 EEF2 HNRNPA2B1 ATP5F1E ARL6IP5 SLC25A5 MYL12B HNRNPA1 HCST MYH9 MSN EEF1G BTG2 TRAPPC1 NACA HMGN2 GNG2 MYL12A RGL4 LMO4 CRIP1 PNRC1 PPIB EEF1D SUMO2 NEAT1 FOS SARAF ETS1 TNF SNRPD2 RPS5 STK17B RPL36A ARPC2 RHOB SEPTIN1 SAT1 CYTIP EVL GUK1 MIF JUND CFL1 COX6B1 C12ORF57.\n",
      "Cell 2:\n",
      "RPLP1 CCL4 TMSB4X B2M RPL41 ACTB MT-CO1 MALAT1 RPL13 RPS19 RPL28 EEF1A1 RPS12 CCL5 HLA-B MT-CO2 TPT1 RPS27 MT-CO3 RPL32 MT-ND3 HSPA1A XCL1 RPS4X KLF6 DNAJB1 RPS27A RPS3 RPL10 RPS15A RPLP0 RPS18 RPS28 RPS23 MT-ATP8 RPL11 RPL30 FOS RPS2 RPS7 NR4A2 RPS14 RPS6 HLA-C RPL39 RPL18A RPL37 RPL19 JUN SH3BGRL3 TMSB10 MT-CYB RPL34 PFN1 RPL26 RPL15 HSP90AA1 TRBV9 RPS15 RPS25 RPL29 RPL17 RPS24 VIM RPS29 MT-ND4L ACTG1 RPL7A RPL36 RPL18 RPS21 RPS3A RPS8 IFITM1 H3-3B HLA-A PPP1R15A PTMA MT-ATP6 FTH1 BTG1 GAPDH RPL23A RPL3 ATF3 CD74 DUSP1 UBA52 RPL35 PNRC1 RPS16 RPL37A ATP5F1E RACK1 RPL21 RPLP2 RPL9 RPL12 RPL10A HSPA8 RPS13 EIF1 RPL8 MT-ND1 SERF2 RPL38 TNFAIP3 ARHGDIB SRGN RPL36A RPSA RPS5 IFITM2 SPRY1 RPL6 PPIA CFL1 FAU CD7 EEF2 RPL35A HCST MT-ND5 SLA IL32 CORO1A NACA RPL24 HSPA1B CRIP1 HNRNPA2B1 H3-3A CYTIP RPL27 ZFP36L2 RPL22 CD52 MT-ND4 TRAV19 ZFP36L1 CKLF RPS9 PPIB RPL14 EVL COX6B1 MYL6 RPL4 ZFP36 UBC RHOB FTL IL7R EEF1B2 UBB RPL27A TOMM7 GYPC EEF1G RAC2 CAPZB OAZ1 MAP3K8 PSMB10 HSPE1 RPL5 CD3E SARAF CD99 PABPC1 SRP14 RBM3 CHCHD2 CD3D CDKN1A NDUFS5 MYL12B EIF3F TAGAP RHOH PFDN5 SRSF7 RPL31 NR4A1 MCL1 MYL12A OST4 HMGB1 RAB9A S100A10 SAMSN1 NKG7 FYB1 S100A4 NFKBIA CNBP MT-ND2 PLAAT4 DYNLL1 SLC2A3.\n",
      "Cell 3:\n",
      "CCL4 MALAT1 CCL4L2 B2M RPLP1 MT-CO2 EEF1A1 MT-CO1 CCL5 JUN TMSB4X IFNG FOS MT-CO3 RPL41 MT-ND4L DNAJB1 HSP90AA1 RPS12 MT-CYB ACTB RPL28 KLF6 HSPA1A PTMA TPT1 RPS19 RPS4X MT-ND3 CD69 RPL30 RPS27 RPS27A RPS28 RPL13 HLA-B RPL10 RPL32 RPS15A MT-ATP8 RPLP2 RPS3 RPL19 RPS18 RPL37 RPL18A RPS21 RPS3A RPL18 RPS14 PPP1R15A ANXA1 MT-ND1 H3-3B RPS15 RPL34 RPL39 SH3BGRL3 MT-ATP6 NR4A1 RPS8 RPL35A MT-ND2 RPS25 RPS23 ATF3 DUSP1 GAPDH RPL7A RPL26 RPL8 RPS13 RPL29 MT-ND4 RPS7 TMSB10 HLA-E IL32 ACTG1 RPL12 TRBV7-9 HSPA1B RPL14 CD8A UBB RPL3 RPS29 RPS24 RPS16 RPL36A HLA-C RPL5 HSPA8 S100A4 RPL22 RPS2 TMA7 NKG7 TNFAIP3 RPL36 RPL21 RPS5 MT-ND5 RPLP0 ATP5F1E RPL11 HSP90AB1 FAU TNF CKLF RPS6 JUND RPL9 FOSB FYB1 MYL6 FTH1 CD3E BTG1 RPL6 MYO1E ZFP36 PTPRCAP KLRB1 MCL1 VIM HCST RPL37A CD7 RPL4 FTL COX4I1 SRGN EIF1 RPS9 IFITM2 PCBP2 ITM2B NFKBIZ EEF1D SARAF RPL13A CDKN1A NFKBIA PNRC1 CFL1 MYL12A OAZ1 UBA52 IFITM3 TAGAP LY6E RPL24 ARHGDIB LDLRAD4 RPL35 TSC22D3 GZMA CRIP1 JUNB CD3D PFN1 UQCR11 NEU1 NACA RPL27 RPL15 BTF3 PCBP1 ATP1A1 RPL23 HLA-A COX7A2 HNRNPA3 MRPL18 ARGLU1 PABPN1 IFITM1 ZFP36L1 MYADM CORO1A MIF IL2RG SNRPB2 RACK1 HNRNPK CHCHD2 TPM3 ANXA11 HSPA5 SLC25A6 EEF2 CSK CALR RPSA ELF1 RBM39 STK17A COMMD6 CD2.\n",
      "Cell 4:\n",
      "B2M MT-CO1 TMSB4X MT-CO2 RPLP1 HSP90AA1 MALAT1 HSPA8 MT-CO3 EEF1A1 H3-3B HSPA1A RPL41 MT-CYB ACTB RPS27 HSPA1B RPS18 KLF6 RPL28 RPS19 HLA-B MT-ATP6 RPL10 RPL13 MT-ATP8 RPS12 RPS28 RPL11 TPT1 MT-ND3 RPS27A CCL5 JUNB RPL39 DNAJB1 RPS2 UBC ANXA1 RPS21 MT-ND4 RPL32 MT-ND4L RPS15 RPS3 ATF3 RPL34 DUSP1 RPL21 TNFAIP3 RPLP2 RPS3A RPL37 RPL30 RPS15A BTG1 PTMA MT-ND1 RPS23 PFN1 NR4A2 RPS4X RPS14 TRBV4-1 HLA-C RPS6 HLA-A RPL12 RPL3 RPL29 PNRC1 SARAF RPL19 RPL18A RPL36 VIM RPS8 RPS7 NR4A1 TMSB10 RPSA RPL26 MT-ND5 RPS29 RPS13 RPL15 EIF1 RPL17 RPL35A CSF1 FAU GAPDH HSP90AB1 MYL12A S100A4 RPL14 FTL CD69 ACTG1 RPL18 SH3BGRL3 RPL7A CD8A ADGRE5 GZMA JUN PLAAT4 PTPRCAP NEAT1 RPLP0 RPS25 RPL22 DNAJA1 CD7 RPS24 SELENOK RPL24 RPS9 RPL5 RPS5 FOS LGALS1 MT-ND2 RPL37A STK17B NACA EEF1B2 SAT1 CYTIP RPL23A JUND CORO1A RPL38 RPL13A RPL8 PPIA FYB1 CFL1 SRGN ZFP36L2 RHOH EEF1G RPS16 HSPE1 EEF2 IL32 KLRB1 RPL9 HSPH1 TSC22D3 UBA52 SEPTIN6 PFDN5 RPL10A CD3G YBX1 CD3E HLA-E RPL4 ARHGDIB PTGER4 LIMD2 CACYBP HSPB1 SERF2 FOSB HNRNPA0 ITM2B CKLF FERMT3 RACK1 MCL1 CRIP1 RPL35 PTGES3 CAPZB RPS26 SDCBP TPI1 RAC1 CAP1 COX7C RPL36A FRMD4B TOMM7 PPP1R15A IFITM1 HSPD1 GNB2 IL7R PMAIP1 ARPC3 ATP5F1E CD3D MTRNR2L12 LSP1 MYL6 CYCS TIPARP SPATS2L.\n",
      "Cell 5:\n",
      "MT-CO2 GNLY MT-CO3 MT-ND3 MT-CYB MT-CO1 RPLP1 MALAT1 TMSB4X MT-ATP6 EEF1A1 RPL41 MT-ATP8 HSP90AA1 RPS19 MT-ND1 RPS27 MT-ND4L GZMA CCL5 RPL13 MT-ND2 RPL30 RPL28 RPS23 RPS18 RPS27A RPL10 RPS15A ACTB B2M MT-ND4 RPS4X RPS12 RPS3 HSPA1A PTMA TPT1 RPL32 RPS8 RPS14 RPS6 RPL26 RPL34 RPL39 RPS28 GAPDH CD7 RPS24 RPL18A RPS3A RPL11 RPL8 RPS15 RPL7A RPL3 RPS7 KLF6 RPL35A RPL12 RPL23A RPL19 NKG7 RPL6 RPLP2 RPS9 RPLP0 CD52 RPL29 RPL37 TMSB10 RPL37A RPS13 RPL35 DNAJB1 MT-ND5 RPS25 RPS2 RPL15 FTL TYROBP RPS21 EEF1B2 HSPA1B RPL36A HSPA6 UBC ACTG1 MTRNR2L12 ZFP36L1 RPL18 FAU RPL22 RPL9 RPS29 NACA FTH1 NR4A1 RPL21 RPL17 ATP5F1E RPS5 RACK1 PABPC1 RPL10A IER5 CORO1A RPL14 RPL36 PFN1 H3-3B SRGN TAGAP MYL12A IFITM2 RPS16 RPL27 HCST VIM HLA-B HOPX CFL1 ATF3 SERF2 SH3BGRL3 NPM1 HLA-A EEF1D NR4A3 IFNG RPL27A CKLF EEF1G UBB RPL4 H2AJ HSP90AB1 EIF1 YBX1 PLAAT4 TRDC RPS26 BTG1 GEM CD3E COTL1 HSPA8 RPL24 FCER1G PNRC1 ARHGDIB RPL5 TOMM7 IFITM1 RPL38 RAC2 CYBA TRA2B MIF SEC11A CD247 RPL7 RPL31 OAZ1 PFDN5 CD69 KRT86 RPL13A RPSA EEF2 BTF3 COX7C MYL6 RHOC UBXN4 UBA52 DBI HMGN2 GSTK1 HSPB1 KLRB1 KIR2DL4 HLA-E FOSB CIRBP PTPN22 PHLDA1 TMA7 SYTL3 ATP5MG RHOA GNAS GYPC HLA-C H3-3A COX6C KLRC2 PARK7 NFKBIA ATP5MC2.\n",
      ".\n",
      "Abstract summary:\n"
     ]
    }
   ],
   "source": [
    "print(formatted_ds[0]['model_input'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07376fd5",
   "metadata": {},
   "source": [
    "# Load C2S Model\n",
    "\n",
    "Now, we will load a pretrained C2S model that has been trained on a large corpus of single-cell data and biological text. This model has learned to understand the \"language\" of cells and can be used for various interpretation tasks. We will use the `CSModel` class from the `cell2sentence` library to load the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "0317ac96",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda\n"
     ]
    }
   ],
   "source": [
    "# Define the path to your pretrained model and a directory to save model-related files\n",
    "model_path = \"vandijklab/C2S-Scale-Pythia-1b-pt\"\n",
    "save_dir = \"/home/sr2464/scratch/C2S_API_Testing/Cache_Dir\"\n",
    "save_name = \"natural_language_interpretation_1B_model\"\n",
    "\n",
    "# Initialize the CSModel object\n",
    "csmodel = cs.CSModel(\n",
    "    model_name_or_path=model_path,\n",
    "    save_dir=save_dir,\n",
    "    save_name=save_name\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "7fd868a0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cuda'"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "98fb6e3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "7620bdf2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPTNeoXForCausalLM(\n",
       "  (gpt_neox): GPTNeoXModel(\n",
       "    (embed_in): Embedding(50304, 2048)\n",
       "    (emb_dropout): Dropout(p=0.0, inplace=False)\n",
       "    (layers): ModuleList(\n",
       "      (0-15): 16 x GPTNeoXLayer(\n",
       "        (input_layernorm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "        (post_attention_layernorm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "        (post_attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (post_mlp_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (attention): GPTNeoXSdpaAttention(\n",
       "          (rotary_emb): GPTNeoXRotaryEmbedding()\n",
       "          (query_key_value): Linear(in_features=2048, out_features=6144, bias=True)\n",
       "          (dense): Linear(in_features=2048, out_features=2048, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (mlp): GPTNeoXMLP(\n",
       "          (dense_h_to_4h): Linear(in_features=2048, out_features=8192, bias=True)\n",
       "          (dense_4h_to_h): Linear(in_features=8192, out_features=2048, bias=True)\n",
       "          (act): GELUActivation()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_layer_norm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "    (rotary_emb): GPTNeoXRotaryEmbedding()\n",
       "  )\n",
       "  (embed_out): Linear(in_features=2048, out_features=50304, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    os.path.join(save_dir, save_name),\n",
    "    cache_dir=os.path.join(save_dir, \".cache\"),\n",
    "    trust_remote_code=True\n",
    ").to(device)\n",
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23d5cbde",
   "metadata": {},
   "source": [
    "# Generate Natural Language Interpretations of Cell Subsets\n",
    "\n",
    "With our model loaded, we can now ask it to interpret different subsets of cells from our dataset.\n",
    "\n",
    "Here, we select a sample of cells from our dataset to interpret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "475cb926",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'multi_cell_groupings': [9744, 10316, 10819, 9040, 8557],\n",
       " 'batch_label': 'A29',\n",
       " 'tissue': 'mesenteric lymph node',\n",
       " 'organism': 'Homo sapiens',\n",
       " 'abstract': ''}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "multi_cell_arrow_ds[10000]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d76d4d2",
   "metadata": {},
   "source": [
    "Note that the Dominguez-Conde dataset contains immune cells taken from many different tissues around the body. These are a group of 5 immune cells from a mesenteric lymph node tissue of one of the patients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1805db5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Input Prompt ---\n",
      "Examine the top 200 genes ordered by descending expression level in 5 Homo sapiens cells. Based on this information, generate an abstract summary which summarizes the biological insights contained in these cells.\n",
      "Cell 1:\n",
      "CD74 EEF1A1 MT-ND3 HLA-DRA RPS27 IGHV5-51 RPLP1 MT-CO2 RPL41 FOS RPL13 RPS19 MT-CO1 RPL32 RPL10 RPS12 RPS8 RPL30 ACTB MT-CYB HLA-DRB1 CXCR4 B2M MALAT1 RPS15A RPL8 RPL28 MT-ATP8 RPL35A RPS3A RPL18A MT-ND4L DUSP1 RPS4X TMSB4X RPL39 MT-CO3 TPT1 RPS27A RPS15 RPS28 RPS3 RPS21 RPS13 RACK1 RPL37 RPS18 RPL35 MT-ATP6 RPL26 RPS25 RPL29 CD69 RPL21 RPL19 RPS2 FAU RPL23A RPL34 RPL15 RPS5 BTG1 MT-ND4 CD52 CD37 RPL36 RPS14 RPS9 MTRNR2L12 CD79A RPLP0 JUN HLA-B RPS6 HLA-DPA1 EIF1 RPS23 RPSA RPS7 RPL27 RPL11 TCL1A HLA-DRB5 H3-3B SLC2A3 TXNIP RPS24 RPS29 RPL37A IGKC RPLP2 RPL7A CD83 CYBA PABPC1 RPL18 HLA-DQB1 MT-ND5 RPL9 RPL38 RPL10A RPL5 MT-ND1 HLA-A FTL RPL6 HLA-E EEF2 PTMA TSC22D3 FTH1 TAGAP UBC MS4A1 TMSB10 RPL36A PCBP1 LY9 ASH1L SARAF UBA52 ZFP36L1 NACA RABAC1 EIF4A3 HLA-DPB1 JUND POU2F2 RPS10 MZT2B PCBP2 RPL7 SESTD1 HLA-DQA2 RPL14 RPL13A POLD4 ECHDC1 PFDN5 RBM38 UQCRB ARHGDIB HLA-C PFDN1 HECTD1 IDI1 GNB5 TUBA1A SOCS3 SMAP2 DCAF7 LINC02397 SNHG8 IER2 HMGB1 TNFRSF13C ENO1 ACTG1 MT-ND2 COX7C HNRNPA1 BAX FCER2 RPL24 PKM SRGN PNPLA2 ARPC2 CHCHD10 DDX50 COX19 HVCN1 NCF1 CCNI FCMR HTT PTPRCAP GOLM2 JUNB CD79B SLC44A2 PFN1 TMEM258 MAX RPL22 REL SSR1 LAMTOR4 CD19 ORAI2 OAZ1 CHI3L2 KLF6 LIG1 LTB H2BC4 KLF2 HNRNPA2B1 ZNF106 KLF4.\n",
      "Cell 2:\n",
      "MALAT1 EEF1A1 RPS27 RPS12 RPL41 RPL13 RPLP1 RPL39 RPS14 TPT1 RPS19 RPL37 RPL30 RPL11 RPL32 RPS15A RPL10 RPS27A RPS18 RPS29 RPS8 RPS23 RPL34 RPL18A RPL28 RPS28 RPS3A TMSB4X RPS4X RPL19 RPS6 RPL26 RPS2 RPS21 RPL29 RPS25 RPS3 RPL21 RPS13 MT-ND3 RPLP2 RPL3 RPL12 B2M ACTB RPL35A MT-CO2 MT-CO1 RPL22 RPL18 MT-CO3 RPL7A RPL5 RPS9 RPL23A MT-ND4L RPSA RPL8 RPLP0 RPL17 RPL15 RPS24 RPS15 RPL37A TMSB10 MT-ATP8 MT-ATP6 RPL36 JUNB HLA-B FAU MT-CYB RPL14 RPL13A RPS16 EEF1B2 RPS7 FTH1 RPL6 RPL9 RPS5 MTRNR2L12 RPL10A RPL4 PTMA ACTG1 RPL24 NACA EEF1G S100A4 RPL38 RPL27 RPS10 TRBC1 SARAF EEF1D HLA-A TMA7 RPS11 ATP5MC2 PPIA PABPC1 MT-ND1 ZFP36 DDX5 SLC2A3 UBA52 PFN1 SLC25A3 UQCRB FYB1 HLA-E IFITM1 FLT3LG GAPDH SNHG29 ZFP36L2 VIM RPS20 NOSIP SH3BGRL3 IL32 RPL35 RPL23 EIF1 COX4I1 ATP5ME FTL SERF2 BTG1 ATP6V1G1 RACK1 CD3G CCR7 YBX1 EEF2 RAC2 PSME1 ARHGDIB CORO1A MT-ND5 MYL6 EIF3K GPR183 CD3D ITM2B SUMO2 TUBB GAS5 UXT HLA-C PCBP2 EIF3H ATP5F1A SOD1 PTPRCAP CD7 RP11-290D2 RPL31 FUS HINT1 PRKCQ-AS1 OST4 C12ORF57 H3-3A ACTR1A HSPA8 CCNH LITAF DNAJC8 GHITM RYK RPL36A STK4 CFL1 ATP5F1E TRBV6-1 RPL27A PSME2 TRAV12-1 LDHB NEDD8 KLF2 ARPC5 ARPC3 PCBP1 COMMD6 EVL TCF7 MSL3 ARL4D ZNF330 ANP32B LIMD2 CD55 TBCA ESYT1 NSMCE1 SNRPC DAD1.\n",
      "Cell 3:\n",
      "EEF1A1 MALAT1 B2M RPL41 RPS27 TMSB4X MT-ND3 RPS28 RPL30 HLA-B RPS15A MT-CYB DUSP1 RPL32 RPL39 RPL37 RPL28 RPL13 CD69 MT-CO1 RPL10 RPLP1 RPS3 TRBV5-1 RPS12 SLC2A3 HLA-C TPT1 RPS23 RPL26 RPS27A RPL34 FOS RPS18 RPL9 RPL35A RPS29 RPL11 MT-ND4L RPL19 RPS21 RPS15 MT-ATP8 RPL29 RPS6 MT-CO2 TMSB10 ACTG1 MT-CO3 RPS7 RPL12 TRAV14DV4 JUNB RPS2 ACTB RPS8 RPL8 RPS3A RPL18A HLA-E RPL3 RPS14 MT-ATP6 EIF1 RPL21 RPL14 RPL18 H3-3B MTRNR2L12 KLF6 EEF1D JUN MT-ND2 RPS19 MT-ND1 MT-ND4 RPL37A RPSA RPL15 RPL7A KLRB1 COX7C RPL36AL RPS4X RPL35 TSC22D3 RPS25 FAU RPS13 IL32 RPLP2 RPL5 ZFP36L2 HLA-A RPL36 HSP90AA1 DDIT4 EEF1G TCF7 HSP90AB1 SEPTIN1 RPS24 TUBA4A DNAJB1 PABPC1 RPLP0 BTG1 RPL24 RPS10 EIF5B CD5 DDX5 NEAT1 CYTIP ZFP36L1 UBB LAPTM5 RPS20 FTH1 CD3D RACK1 SRGN FTL IFITM1 BTF3 FXYD5 HNRNPH2 YWHAZ NACA SAP18 DNAJA1 PFN1 CD44 SRSF7 BTG2 RPL38 CFL1 EEF1B2 ICOS CTLA4 HNRNPA2B1 BAZ2A ARHGDIB WIPF1 PSMB8 HSPA1A SON SS18L2 RPL17 GAS5 ZNHIT1 RPL10A EVI2B RPS9 ERLEC1 PRR13 RBMX PFDN5 MT-ND5 PARP1 HLA-F RPL7 RPL23A RPS5 FCMR CMTM7 RPS16 PCBP2 SARAF FYB1 RAN CCNI CHCHD2 MYL6 RPL4 IFITM2 GAPDH RPL22 COX20 CD3E CD2 CCR4 COX14 MAGOH ARAP2 MRPL23 UFD1 LAT RPL6 VIM UBA52 CREM NPM1 SPN DAD1 OAZ1 TRBC2 DDX24 RBM39 EIF3F.\n",
      "Cell 4:\n",
      "EEF1A1 RPS27 RPL10 RPL13 RPS12 RPL30 MALAT1 RPLP1 RPL41 RPL32 MT-ND3 MT-CO2 RPS15A RPS28 RPS23 RPS3 RPL28 RPS27A RPS3A RPS8 RPL18 MTRNR2L12 RPL29 RPL19 RPS21 TPT1 B2M MT-CO3 RPL34 RPS4X RPL26 RPS19 RPS13 RPL11 RPS14 RPS18 RPL37 RPLP0 RPS2 RPL39 RPL12 RPL3 RPL14 RPL5 MT-ATP8 RPS25 RPL35A RPL21 RPS6 RPS15 RPS24 MT-CO1 RPS29 MT-ND4L TMSB4X RACK1 TMSB10 RPL6 RPL7A RPS9 RPL9 ACTB RPL22 RPL15 MT-CYB RPL8 PABPC1 RPL10A JUNB RPL36 FOS RPSA DNAJB1 TSC22D3 RPL37A RPL38 RPS5 MT-ATP6 VIM RPL18A PTMA RPLP2 EEF2 FAU NACA HLA-C RPL36A RPS16 EIF1 RPL23A NPM1 RPL27 RPL24 RPS7 RPL4 DUSP1 RPL35 RPL13A HLA-B UBA52 ACTG1 HNRNPA1 COX4I1 MT-ND2 EEF1G RPL17 BTG1 HNRNPDL EEF1D MT-ND1 PFN1 FTH1 CD69 KLF2 ATP5MC2 SELL PNRC1 BTF3 GAS5 ARHGDIB MT-ND4 SLC2A3 EEF1B2 SNHG29 HLA-A LTB H3-3A RPL23 FTL SRSF7 RPL31 FOSB LINC02446 TRBV6-1 IFITM2 PTGER4 PPP1R15A RPL36AL ZFP36 HCST EIF3E PTPRC SARAF CXCR4 RHOA SLC25A3 ZNF331 R3HDM2 EIF3G PFDN5 LDHB S100A6 CD48 ATP5F1E CD8A CD3E PRR13 HSPA8 RPL27A EIF3H TMA7 TCF7 LCP2 SNHG3 NR4A1 HNRNPA2B1 CD3D ATP5ME TAX1BP1 COX7C SNHG8 PPIA HLA-E PCED1B-AS1 TRAF3IP3 CD55 SLC9A3R1 MYL12A RPL7 SLC27A5 DUSP2 RPS26 H3-3B ST13 HLA-F TXNIP WDR1 PICALM LIMD2 MT-ND5 CD52 GYPC NFKBIA GZMM HSP90AA1 TSTD1 YBX1 PCBP2 CAPZB TOMM20.\n",
      "Cell 5:\n",
      "MTRNR2L12 RPS27 EEF1A1 RPL41 RPL13 RPS12 RPL10 RPL30 TMSB4X RPL32 MALAT1 RPLP1 JUNB RPS28 RPL39 RPL19 RPL34 RPL18A RPL28 RPS14 RPL11 RPS3A RPS15A MT-CO2 ACTB RPS23 RPL35A RPL37 TPT1 RPS18 MT-ND3 MT-CYB RPS4X RPS8 RPS3 RPL29 RPS21 RACK1 RPL12 RPS29 RPL3 RPS15 RPS2 RPS27A PTMA RPLP0 RPL7A RPL5 RPSA MT-ND4L B2M MT-ATP8 RPL26 RPS24 NACA ACTG1 MT-CO1 RPL15 RPL14 RPL23A RPS19 RPL9 FAU RPL36 RPS6 MT-CO3 RPLP2 RPS25 RPS13 RPS9 HLA-B RPS16 RPL18 RPL8 RPL6 RPS7 MT-ATP6 RPL17 IFITM1 RPS5 HLA-A RPL38 RPL7 RPL21 ZNF331 FTL RPL37A MT-ND4 MT-ND1 EIF1 FOS EEF2 RPL10A BTG1 FTH1 RPL31 SARAF NPM1 RPL4 DNAJB1 TRBV14 RPL36A RPL22 PTPRCAP TRBV6-6 CD37 HLA-C PABPC1 CYTIP LDHB ARHGDIB EEF1G PCBP2 MYL12A SNHG6 RPL27A RPL35 PPIA TMSB10 RPL27 ZFP36 NOP53 CD69 RBM39 BTF3 SUMO2 MT-ND2 GPSM3 SNHG29 HINT1 COMMD6 CCNI RPS17 ATP5F1E EEF1D CORO1A PFN1 HSPA8 DUSP2 HNRNPH1 PPP1R2 RAC2 PFDN5 CFL1 CAP1 HLA-E H3-3B STK17B MIF SSH2 SON GABARAP PNRC1 MYH9 RPL24 FAM177A1 COX4I1 HSPE1 EIF3K RHOA TOMM7 BTG2 RGS10 LEPROTL1 PTPRC SELL DDX39B FXYD5 ARPC1B TXNIP SRSF3 SNRPD2 TAF15 RIPOR2 ZFAS1 HNRNPUL1 MPG CD3E VAMP2 EEF1B2 SATB1 SLC25A6 GAPDH CYCS TRAM1 SLC25A3 MTRNR2L8 TMC8 PCBP1 RPS26 H3-3A ITM2B TMA7 SKP1 LSM4 EDF1 EIF5A GIMAP7 ARPC3 YWHAZ.\n",
      ".\n",
      "Abstract summary:\n"
     ]
    }
   ],
   "source": [
    "# Format prompts for both cell subsets\n",
    "input_prompt = formatted_ds[10000]['model_input']\n",
    "\n",
    "print(\"--- Input Prompt ---\")\n",
    "print(input_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "016d6924",
   "metadata": {},
   "source": [
    "Finally, we will use our loaded C2S model to generate natural language summaries for each of our cell subsets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "0c478051",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate interpretation for the ileum cells\n",
    "natural_language_interpretation = csmodel.generate_from_prompt(\n",
    "    model=model,\n",
    "    prompt=input_prompt,\n",
    "    do_sample=True,\n",
    "    max_num_tokens=1024,\n",
    "    temperature=0.7,\n",
    "    top_k=30,\n",
    "    top_p=0.9\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "c1cfe9cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Interpretation for cells ---\n",
      "The study used single-cell sequencing to analyze immune cells in 16 human tissues from 12 donors, generating a dataset of ~360,000 cells. They developed CellTypist, a machine learning tool for precise cell type annotation, and identified tissue-specific features and clonal architecture of T and B cells. This approach provides a foundation for identifying highly resolved immune cell types using a common reference dataset and antigen receptor sequencing..\n"
     ]
    }
   ],
   "source": [
    "print(\"--- Interpretation for cells ---\")\n",
    "print(natural_language_interpretation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1ae0d82",
   "metadata": {},
   "source": [
    "The summary produced by the model correctly identified some details about the dataset it was sampled from, since this dataset (immune cells from Dominguez-Conde et al.) was a part of the C2S pretraining dataset. Sampling with different temperature and on different datasets can give different interpretations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e55401f",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "\n",
    "In this tutorial, we have seen how a pretrained Cell2Sentence model can be used to generate natural language interpretations of single-cell datasets. By providing the model with a collection of cell sentences, it can produce a high-level summary of the biological information contained within those cells. This capability is a powerful tool for researchers to quickly gain insights into their data and understand the biological context of cell clusters in an accessible way.\n",
    "\n",
    "The ability to interpret complex single-cell data in natural language opens up new avenues for data exploration and hypothesis generation. As C2S models become more sophisticated, we can expect them to provide even more detailed and nuanced interpretations, further bridging the gap between computational analysis and biological understanding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75094e42",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cell2sentence2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
